#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
参数更新工具

供Agent调用的参数更新工具，用于在对话过程中更新洪水预报参数
"""

from langchain.tools import BaseTool
from typing import Dict, Any, Optional, Type
from pydantic import BaseModel, Field
import json
from loguru import logger


class ParameterUpdateInput(BaseModel):
    """参数更新工具的输入模型"""
    reservoir_name: Optional[str] = Field(None, description="水库名称")
    forecast_model: Optional[str] = Field(None, description="预报模型")
    start_time: Optional[str] = Field(None, description="开始时间，格式：YYYY-MM-DD HH:MM")
    forecast_duration: Optional[int] = Field(None, description="预报时长（小时）")
    rainfall_forecast_period: Optional[int] = Field(None, description="降雨预见期（小时）")
    target_variable: Optional[str] = Field(None, description="目标变量")


class ParameterUpdateTool(BaseTool):
    """参数更新工具
    
    允许Agent在对话过程中更新洪水预报参数
    """
    
    name: str = "update_parameters"
    description: str = """
    更新洪水预报参数。当用户提供了新的参数信息时使用此工具。
    
    参数说明：
    - reservoir_name: 水库名称（如：大伙房水库、三峡水库等）
    - forecast_model: 预报模型（如：新安江模型、LSTM模型等）
    - start_time: 开始时间，格式：YYYY-MM-DD HH:MM
    - forecast_duration: 预报时长，单位小时（如：24、48、72）
    - rainfall_forecast_period: 降雨预见期，单位小时（如：12、24、48）
    - target_variable: 目标变量（如：流量、水位、库容）
    
    只传入用户明确提到的参数，未提及的参数不要传入。
    """
    args_schema: Type[BaseModel] = ParameterUpdateInput
    agent_instance: Optional[Any] = Field(default=None, exclude=True)
    
    def __init__(self, agent_instance=None, **kwargs):
        super().__init__(**kwargs)
        object.__setattr__(self, 'agent_instance', agent_instance)
    
    def _run(self, *args, **kwargs) -> str:
        """执行参数更新，兼容位置参数和关键字参数
        
        Args:
            *args: 可能包含参数字典
            **kwargs: 参数字典，包含要更新的参数和LangChain的回调管理器
        Returns:
            str: 更新结果描述
        """
        try:
            if not self.agent_instance:
                return "错误：未设置代理实例"
            
            # 添加调试日志（使用INFO级别以便在前端显示）
            logger.info(f"参数更新工具接收到的args: {args}")
            logger.info(f"参数更新工具接收到的kwargs: {kwargs}")
            
            # 处理LangChain传入的参数格式
            if args:
                # 如果第一个参数是字符串（JSON格式），尝试解析
                if isinstance(args[0], str):
                    try:
                        # 尝试直接解析JSON
                        parsed_params = json.loads(args[0])
                        if isinstance(parsed_params, dict):
                            kwargs.update(parsed_params)
                            logger.info(f"解析JSON参数成功: {parsed_params}")
                    except json.JSONDecodeError as e:
                        # 如果直接解析失败，尝试提取第一个JSON对象
                        try:
                            # 查找第一个完整的JSON对象
                            text = args[0].strip()
                            if text.startswith('{'):
                                # 找到第一个JSON对象的结束位置
                                brace_count = 0
                                end_pos = 0
                                for i, char in enumerate(text):
                                    if char == '{':
                                        brace_count += 1
                                    elif char == '}':
                                        brace_count -= 1
                                        if brace_count == 0:
                                            end_pos = i + 1
                                            break
                                
                                if end_pos > 0:
                                    json_part = text[:end_pos]
                                    parsed_params = json.loads(json_part)
                                    if isinstance(parsed_params, dict):
                                        kwargs.update(parsed_params)
                                        logger.info(f"提取JSON参数成功: {parsed_params}")
                                else:
                                    logger.error(f"无法找到完整的JSON对象: {text[:100]}...")
                            else:
                                logger.error(f"参数不是JSON格式: {text[:100]}...")
                        except Exception as extract_error:
                            logger.error(f"JSON参数提取失败: {extract_error}")
                # 如果第一个参数是字典，直接合并
                elif isinstance(args[0], dict):
                    kwargs.update(args[0])
                    logger.info(f"合并字典参数: {args[0]}")
            
            logger.info(f"处理后的kwargs: {kwargs}")
            
            # 过滤掉LangChain内部参数和None值
            updates = {k: v for k, v in kwargs.items() 
                      if k not in ['run_manager'] and v is not None}
            
            logger.info(f"过滤后的updates: {updates}")
            
            # 验证参数有效性
            validation_result = self._validate_parameters(updates)
            if not validation_result['valid']:
                return f"参数验证失败：{validation_result['message']}"
            
            if not updates:
                return "没有需要更新的参数"
            updated_params = []
            for key, value in updates.items():
                if key in ['reservoir_name', 'forecast_model', 'start_time', 
                          'forecast_duration', 'rainfall_forecast_period', 'target_variable']:
                    self.agent_instance.collected_parameters[key] = value
                    updated_params.append(f"{key}: {value}")
                    logger.info(f"更新参数 {key} = {value}")
            if updated_params:
                result = f"已更新参数：{', '.join(updated_params)}"
                
                # 如果用户刚提供了水库名称，自动生成其他参数的推荐值
                if 'reservoir_name' in updates and len(self.agent_instance.collected_parameters) == 1:
                    recommendations = self.agent_instance._generate_default_recommendations()
                    
                    # 更新推荐参数
                    for key, value in recommendations.items():
                        self.agent_instance.collected_parameters[key] = value
                        logger.info(f"自动推荐参数 {key} = {value}")
                    
                    # 更新阶段为推荐确认
                    self.agent_instance.current_stage = "recommendation"
                    
                    result += f"\n\n为您推荐以下参数：\n"
                    for key, value in recommendations.items():
                        param_names = {
                            'forecast_model': '预报模型',
                            'start_time': '开始时间', 
                            'forecast_duration': '预报时长（小时）',
                            'rainfall_forecast_period': '降雨预见期（小时）',
                            'target_variable': '目标变量'
                        }
                        result += f"- {param_names.get(key, key)}: {value}\n"
                    
                    result += "\n这些推荐参数您确定吗？如果不确定，请告诉我需要修改哪个参数。"
                
                return result
            else:
                return "没有需要更新的参数"
        except Exception as e:
            logger.error(f"参数更新失败: {e}")
            return f"参数更新失败: {str(e)}"
    
    def _validate_parameters(self, updates: Dict[str, Any]) -> Dict[str, Any]:
        """验证参数有效性
        
        Args:
            updates: 要验证的参数字典
            
        Returns:
            Dict: 包含验证结果的字典，格式为 {'valid': bool, 'message': str}
        """
        try:
            for key, value in updates.items():
                if key == 'reservoir_name':
                    if not isinstance(value, str) or len(value.strip()) == 0:
                        return {'valid': False, 'message': '水库名称不能为空'}
                    if len(value) > 50:
                        return {'valid': False, 'message': '水库名称长度不能超过50个字符'}
                
                elif key == 'forecast_model':
                    valid_models = ['新安江模型', 'LSTM模型', 'GRU模型', 'Transformer模型', 'XGBoost模型', 'Random Forest模型']
                    if value not in valid_models:
                        return {'valid': False, 'message': f'预报模型必须是以下之一：{"、".join(valid_models)}'}
                
                elif key == 'start_time':
                    try:
                        from datetime import datetime
                        # 验证时间格式
                        datetime.strptime(value, '%Y-%m-%d %H:%M')
                        # 检查时间是否在合理范围内（不能是过去太久的时间）
                        start_time = datetime.strptime(value, '%Y-%m-%d %H:%M')
                        now = datetime.now()
                        if start_time < now.replace(hour=0, minute=0, second=0, microsecond=0):
                            return {'valid': False, 'message': '开始时间不能早于今天'}
                        if (start_time - now).days > 30:
                            return {'valid': False, 'message': '开始时间不能超过30天后'}
                    except ValueError:
                        return {'valid': False, 'message': '开始时间格式错误，请使用 YYYY-MM-DD HH:MM 格式'}
                
                elif key == 'forecast_duration':
                    if not isinstance(value, int) or value <= 0:
                        return {'valid': False, 'message': '预报时长必须是正整数'}
                    if value > 168:  # 7天
                        return {'valid': False, 'message': '预报时长不能超过168小时（7天）'}
                    if value < 1:
                        return {'valid': False, 'message': '预报时长不能少于1小时'}
                
                elif key == 'rainfall_forecast_period':
                    if not isinstance(value, int) or value <= 0:
                        return {'valid': False, 'message': '降雨预见期必须是正整数'}
                    if value > 72:  # 3天
                        return {'valid': False, 'message': '降雨预见期不能超过72小时（3天）'}
                    if value < 1:
                        return {'valid': False, 'message': '降雨预见期不能少于1小时'}
                
                elif key == 'target_variable':
                    valid_variables = ['流量', '水位', '库容', '入库流量', '出库流量']
                    if value not in valid_variables:
                        return {'valid': False, 'message': f'目标变量必须是以下之一：{"、".join(valid_variables)}'}
            
            return {'valid': True, 'message': '参数验证通过'}
        
        except Exception as e:
            logger.error(f"参数验证过程中发生错误: {e}")
            return {'valid': False, 'message': f'参数验证过程中发生错误: {str(e)}'}
 
    
    async def _arun(self, **kwargs) -> str:
        """异步执行参数更新"""
        return self._run(**kwargs)


# if __name__ == "__main__":
#     # 测试用例
#     print("=== 参数更新工具测试 ===")
    
#     # 创建工具实例（不绑定代理）
#     tool = ParameterUpdateTool()
    
#     # 测试参数更新
#     test_params = {
#         "reservoir_name": "大伙房水库",
#         "forecast_model": "新安江模型",
#         "forecast_duration": 48
#     }
    
#     print(f"测试参数: {test_params}")
#     result = tool._run(**test_params)
#     print(f"更新结果: {result}")